# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/3 18:40 
@Author  : ZhangShenao 
@File    : 7.利用CheckPoint机制实现状态数据的持久化.py 
@Desc    : 利用CheckPoint机制实现状态数据的持久化

在LangGraph中,可以通过CheckPoint机制,实现状态数据的持久化
状态数据与具体的thread_id绑定,不同thread_id对应的状态数据是互相隔离的
CheckPoint机制有多种实现,如PostgresSaver、MongoDB、Redis等
可以将状态数据持久化到不同的存储介质中
"""
import os

import dotenv
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent

from dalle_image_tool import create_dalle_image_tool
from gaode_weather_tool import GaodeWeatherTool
from google_serper_tool import create_google_serper_tool

# 加载环境变量
dotenv.load_dotenv()

# 创建工具列表
tools = [GaodeWeatherTool(), create_dalle_image_tool(), create_google_serper_tool()]

# 创建LLM
llm = ChatOpenAI(model="gpt-4o",
                 temperature=0,
                 openai_api_base=os.getenv("OPENAI_API_BASE"))

# 创建CheckPointer实现
# 这里采用内存实现MemorySaver
checkpointer = MemorySaver()

# 使用预构建的create_react_agent,创建图架构,并设置CheckPointer
# 最终会在Graph的compile方法中,传入CheckPointer
# 将图编译成具有状态持久化功能的可执行组件
agent = create_react_agent(
    model=llm,  # 设置LLM
    tools=tools,  # 设置工具列表
    checkpointer=checkpointer,  # 设置CheckPoint实现
)

# 第一次调用
reply1 = agent.invoke(
    {"messages": [
        ("human", "你好，我是zsa，我喜欢写代码和打篮球"),
    ]},

    # 通过config参数,设置thread_id
    # CheckPointer中的状态数据与具体的thread_id绑定
    # 不同thread_id对应的状态数据是互相隔离的
    config={
        "configurable": {"thread_id": 1}
    }
)
print(f"第1次调用结果: \n{reply1}")

# 第二次调用,使用相同的thread_id
# 此时会从CheckPointer中加载状态数据,即具备了记忆功能
reply2 = agent.invoke(
    {"messages": [
        ("human", "你知道我有什么爱好吗？"),
    ]},

    # 通过config参数,设置thread_id
    # CheckPointer中的状态数据与具体的thread_id绑定
    # 不同thread_id对应的状态数据是互相隔离的
    config={
        "configurable": {"thread_id": 1}
    }
)
print(f"第2次调用结果: \n{reply2}")
